package com.hzw.saas.web.shell.service;

import com.hzw.saas.api.hpc.IHpcAccountProviderService;
import com.hzw.saas.api.hpc.IHpcAccountService;
import com.hzw.saas.api.hpc.IHpcAccountUserService;
import com.hzw.saas.api.hpc.bo.HpcAccountBo;
import com.hzw.saas.api.hpc.bo.HpcAccountProviderBo;
import com.hzw.saas.api.hpc.bo.HpcAccountUserBo;
import com.hzw.saas.api.hpc.enums.HpcProviderType;
import com.hzw.saas.api.property.IPropertyService;
import com.hzw.saas.api.property.enmus.SysDictKeyEnum;
import com.hzw.saas.api.user.IUserAccountService;
import com.hzw.saas.api.user.bo.UserAccountBO;
import com.hzw.saas.api.user.query.UserAccountGetQuery;
import com.hzw.saas.common.config.util.AssertUtil;
import com.hzw.saas.web.shell.event.ShellEvent;
import com.hzw.saas.web.shell.model.ShellChannel;
import com.hzw.saas.web.shell.service.query.HpcAccountQuery;
import com.hzw.saas.web.shell.util.JschEnhanceUtil;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.Session;
import jodd.util.StringUtil;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Service;
import org.springframework.validation.annotation.Validated;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * <p>
 *
 * </p>
 *
 * @author sonam
 * @since 2021/6/3 14:11
 */
@Slf4j
@Service("webShellServiceImpl")
@RequiredArgsConstructor
public class WebShellServiceImpl implements IWebShellService {

    public static Map<String, ShellChannel> SSH_MAP = new ConcurrentHashMap<>();
    public static Map<String, String> SESSION_MAP = new ConcurrentHashMap<>();
    private final static String PING = "#ping#";

    private final IUserAccountService userAccountService;
    private final IHpcAccountUserService hpcAccountUserService;
    private final IHpcAccountService hpcAccountService;
    @Resource(name = IPropertyService.SYS_DICT_BEAN)
    private final IPropertyService<SysDictKeyEnum> propertyService;
    private final IHpcAccountProviderService providerService;

    private final ApplicationEventPublisher publisher;

    private final static ExecutorService executorService = Executors.newCachedThreadPool();

    @Override
    public void connection(@Validated HpcAccountQuery hpcAccountQuery) {
        String destination = hpcAccountQuery.getDestination();
        String userName = hpcAccountQuery.getUserName();
        if (StringUtil.isBlank(hpcAccountQuery.getHpcAccountUserId())) {
            sendToUser(userName, destination, "参数错误");
            return;
        }

        // 获取登录的用户信息
        UserAccountBO account = userAccountService.getAccount(new UserAccountGetQuery().setUserName(userName));
        log.debug("username: {}", userName);
        if (Objects.isNull(account)) {
            log.error("用户不存在: {}", userName);
            sendToUser(userName, destination, "用户不存在");
            return;
        }

        HpcAccountUserBo accountUserBo = hpcAccountUserService.getWithRecord(hpcAccountQuery.getHpcAccountUserId());
        if (Objects.isNull(accountUserBo)) {
            sendToUser(userName, destination, "HPC账号不存在");
            return;
        }
        AssertUtil.assertThrow("用户无此HPC账户权限", !Objects.equals(account.getUserId(), accountUserBo.getUserId()));

        HpcAccountBo hpcAccountBo = hpcAccountService.getById(accountUserBo.getHpcAccountId());
        if (Objects.isNull(hpcAccountBo)) {
            sendToUser(userName, destination, "HPC账号不存在");
            return;
        }

        // 异步执行,与SSH建立连接
        executorService.execute(() -> connectToSSH(hpcAccountBo, hpcAccountQuery));
    }

    @SneakyThrows
    public void connectToSSH(HpcAccountBo hpcAccountBo, HpcAccountQuery hpcAccountQuery) {
        String sessionId = hpcAccountQuery.getSessionId();
        String userName = hpcAccountQuery.getUserName();
        String uuid = hpcAccountQuery.getUuid();

        log.info("connect...");
        Session session = JschEnhanceUtil.getSession(hpcAccountBo.getHost(), hpcAccountBo.getPort(), hpcAccountBo.getHpcAccount(), hpcAccountBo.getHpcPassword(), uuid);
        SESSION_MAP.put(sessionId, uuid);

//        Session session = JschUtil.openSession(hpcAccountBo.getHost(), hpcAccountBo.getPort(), hpcAccountBo.getHpcAccount(), hpcAccountBo.getHpcPassword());
        Channel channel = session.openChannel("shell");

        ShellChannel shellChannel = new ShellChannel();
        shellChannel.setChannel(channel);
        shellChannel.setSession(session);

        //通道连接 超时时间3s
        channel.connect(10000);
        log.info("connect complete...");
        //转发消息，初始消息
        sendToUser(userName, hpcAccountQuery.getDestination(), "连接成功\r");
        //读取终端返回的信息流
        InputStream inputStream = channel.getInputStream();
        shellChannel.setInputStream(inputStream);
        //保存连接
        SSH_MAP.put(sessionId, shellChannel);
        try {
            //循环读取
            byte[] buffer = new byte[1024];
            int i;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            String providerId = hpcAccountBo.getProviderId();
            HpcAccountProviderBo providerBo = providerService.getById(providerId);
            if(Objects.nonNull(providerBo) && Objects.equals(providerBo.getProviderCode(), HpcProviderType.QY.value())) {
                String welcomeMsg = propertyService.getPropertyValue(SysDictKeyEnum.WELCOME_MSG);
//                sendToUser(userName, hpcAccountQuery.getDestination(), StrUtil.format("\033[36;1m{}\033[0m", welcomeMsg));
                sendToUser(userName, hpcAccountQuery.getDestination(), welcomeMsg);
            }
            while ((i = inputStream.read(buffer)) != -1) {
                byte[] payload = Arrays.copyOfRange(buffer, 0, i);
                String msg = new String(payload, "UTF-8");
                sendToUser(userName, hpcAccountQuery.getDestination(), msg);
            }
        } catch (Exception e) {
            log.error("JSCH 连接断开", e);
        } finally {
            log.debug("释放 JSCH 连接");
            //断开连接后关闭会话
           this.closeConnection(sessionId);
            if (inputStream != null) {
                inputStream.close();
            }
        }
    }

    private void sendToUser(String userName, String path, String msg) {
        publisher.publishEvent(new ShellEvent(userName, path, msg));
    }

    @Override
    public void transToSSH(String sessionId, String command) throws IOException {
        if(Objects.equals(command, PING)) {
            command = new String(new byte[]{27}); // 返回ESC
        }

        // 不需要重新连接
        ShellChannel shellChannel = SSH_MAP.get(sessionId);
        if (shellChannel != null) {
            Channel channel = shellChannel.getChannel();
            if (channel != null) {
                OutputStream outputStream = channel.getOutputStream();
                outputStream.write(command.getBytes());
                outputStream.flush();
            }
        }
    }

    @Override
    public void closeConnection(String sessionId) {
        ShellChannel shellChannel = SSH_MAP.get(sessionId);
        try {
            log.debug("close：" + shellChannel);
            int retryCount = 0;
            while(retryCount > 10) {
                if (shellChannel == null) {
                    try {
                        TimeUnit.SECONDS.sleep(1);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    retryCount ++;
                    continue;
                }

                Session session = shellChannel.getSession();
                if (session != null) {
                    session.disconnect();
                }

                InputStream inputStream = shellChannel.getInputStream();
                if (inputStream != null) {
                    try {
                        inputStream.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }

                Channel channel = shellChannel.getChannel();
                if (channel != null) {
                    channel.disconnect();
                }

                break;
            }
        } finally {
            log.debug("remove：{}", sessionId);
            SSH_MAP.remove(sessionId);
            JschEnhanceUtil.closeSession(SESSION_MAP.get(sessionId));
        }

    }
}
